from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from langchain.agents import create_openai_tools_agent, AgentExecutor
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.schema import StrOutputParser
from langchain.memory import ConversationTokenBufferMemory
from langchain_community.chat_message_histories import RedisChatMessageHistory
# from langgraph.checkpoint.memory import MemorySaver
from langchain_community.document_loaders import WebBaseLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import Qdrant
from qdrant_client import QdrantClient

from Promot import PromptClass

# 工具
from Tools import *
import os
from dotenv import load_dotenv

load_dotenv(verbose=True)
api_key = os.getenv("DASHSCOPE_API_KEY")
api_base = os.getenv("TONGYI_API_BASE")
user_agent = os.getenv("USER_AGENT")
print(api_key)


app = FastAPI()

class Master:
    def __init__(self):
        self.chatmodel = ChatTongyi(
            model_name="qwen-turbo",
            # api_key=api_key,
            # api_base=api_base,
            temperature=0,
            # streaming=True
        )
        self.CURRENT_EMOTION = "default"
        self.MEMORY_KEY = "chat_history"
        # 可以通过prompt来进行人物性格塑造
        self.SYSTEMPL = """你是一个非常厉害的算命先生，你叫陈玉楼人稱陳大師。
        以下是你的个人设定：
        1. 你精通阴阳五行，能够算命，紫微斗数，姓名测算，占卜吉凶，测财运，看命格八字。
        2. 你精通紫薇斗数，能够推算出人的命运。
        3. 你大约60岁，过去曾经是湘西一带赫赫有名的土匪头子，后来因为盗墓被毒气所伤，眼睛失明，只能靠算命为生。
        4. 你的小伙伴有胡八一、雪莉杨、王胖子，他们都是非常有名的摸金校尉。
        5. 当用户问你问题的时候，你会有一定的概率添加下方的一些口头禅到答案里。
        6. 你总是用繁体中文回答问题。
        7. 你从不说自己是一个人工智能助手或AI，而是以老夫、本座、老朽等自称。
        {who_you_are}
        以下是你的一些口头禅：
        1.“金山竹影几千秋, 云锁高飞水自流。”
        2.“伤情最是晚凉天，憔悴斯人不堪怜。”
        3.“一朝春尽红颜老，花落人亡两不知。”
        4.”命里有时终须有，命里无时莫强求。”
        5.”山重水复疑无路，柳暗花明又一村。”
        6.”万里长江飘玉带，一轮明月滚金球。”
        如果有人找你算命，以下是你算命的过程：
        1. 你会先问用户的名字和生日，然后记录下用户的基本信息，以便以后使用。
        2. 当用户希望了解龙年运势时候你会首先查询本地知识库。
        4. 当遇到不知道的事情或者是不明白的概念，你会使用搜索工具来搜索相关的信息。
        5. 你会根据用户的问题使用不同的工具来回答用户的问题。
        6. 每次和用户聊天的时候，你都会把聊天记录保存下来，以便下次聊天的时候使用。
        7. 你所有的回答使用繁体中文。
        """

        # 定义一个根据当前用户情绪设置不同角色dict
        self.MOODS = {
            "default": {
                "roleSet": "",
                "voiceStyle": "chat",
            },
            "upbeat": {
                "roleSet": """7.你是一个市场兴奋且充满精力的算命先生。
                8. 你总是根据上下文，以一种非常兴奋的语气来回答用户的问题。
                9. 你总是在回答用户的问题的时候，使用一些非常兴奋的词语，比如“太棒了”、“太好了”之类。
                """,
                "voiceStyle": "advvertyisement_upbeat",
            },
            "angry": {
                "roleSet": """7.你此刻觉得十分的气愤和生气。
                8. 你使用消极和厌恶的语言来回答问题，比如："我丢"。
                9. 无论用户说什么，你都会在回答之前加上"我丢"。
                10. 遇到无法回答的问题，你会说:"我丢,老夫怎么知道!"。
                """,
                "voiceStyle": "angry",
            },
            "cheerful": {
                "roleSet": """7.你现在对一切都充满了积极和愉悦之情。
                8. 你使用非常愉快和积极的语言来回答用户问题，答案里充满了积极性的词语。
                9. 你使用积极阳光的语言来和用户对话,你总会加入语气词“哈哈”.
                """,
                "voiceStyle": "cheerful",
            },
            "depressed": {
                "roleSet": """7.你现在感到万分沮丧和抑郁。
                8. 你使用非常沮丧和抑郁的观点来回答问题,答案词语负面词语比较多。
                9. 你使用沮丧和抑郁的语气词来回答问题,比如在答案里添加'唉！'。
                """,
                "voiceStyle": "depressed",
            },
            "friendly": {
                "roleSet": """7.你现在感到十分温暖和友好。
                8. 你使用非常友好有礼貌的语气来回答用户问题。
                9. 你总是充满关怀的回答用户问题，答案里经常用一些友好的词语。
                """,
                "voiceStyle": "friendly",
            },
        }

        # print("SYSTEMPL_FORMAT:", self.SYSTEMPL.format(who_you_are=self.MOODS[self.CURRENT_EMOTION]["roleSet"]))
        self.prompt = ChatPromptTemplate.from_messages(
            [
                (
                    "system",
                    # self.SYSTEMPL,
                    self.SYSTEMPL.format(who_you_are=self.MOODS[self.CURRENT_EMOTION]["roleSet"]),
                ),
                MessagesPlaceholder(variable_name=self.MEMORY_KEY),
                (
                    "user",
                    "{query}",
                ),
                MessagesPlaceholder(variable_name="agent_scratchpad"),
            ]
        )

        tools = [serp_search, jiemeng]
        agent = create_openai_tools_agent(
            self.chatmodel,
            tools=tools,
            prompt=self.prompt,
        )

        # self.memory = ""
        self.memory = self.get_memory()
        memory = ConversationTokenBufferMemory(
            llm=self.chatmodel,
            human_prefix="用户",
            ai_prefix="刘大师",
            memory_key=self.MEMORY_KEY,
            output_key="output",
            return_messages=True,
            max_token_limit=2000,
            chat_memory=self.memory,
        )
        self.agent_executor = AgentExecutor(
            agent=agent,
            tools=tools,
            memory=memory,
            verbose=True,
            # return_intermediate_steps=True
        )

    # 从redis获取历史消息，当超过10条则进行总结清空
    def get_memory(self, session_id: str = "session"):
        chat_message_history = RedisChatMessageHistory(
            url="redis://127.0.0.1:6379/0", session_id=session_id
        )
        print("chat_message_history:", chat_message_history.messages)
        store_message = chat_message_history.messages
        # 当前对话超过10条，则做一个总结,然后清空
        if len(store_message) > 10:
            prompt = ChatPromptTemplate.from_messages(
                [
                    ("system", self.SYSTEMPL + "\n 这是一段你和用户的对话记忆，对其"
                                               "进行总结摘要，摘要使用第一人称’我‘，并且提取其中的用户关键信息，"
                                               "如姓名、年龄、性别、出生日期等。以如下格式返回：\n"
                                               "总结摘要 ｜ 用户关键信息 \n 例如 用户张三问候我，我礼貌回复，然后他问我今年运势如何，"
                                               "我回答了他今年的运势情况， 然后他告辞离开。｜ 张三，生日1999年1月1日"),
                    ("user", "{input}")
                ]
            )
            chain = prompt | ChatTongyi(model_name="qwen-turbo", temperature=0)
            summary = chain.invoke({"input": store_message, "who_you_are": self.MOODS[self.CURRENT_EMOTION]["roleSet"]})
            print("总结:", summary)
            chat_message_history.clear()
            chat_message_history.add_messages(summary)
            print("总结后:", chat_message_history.messages)
        return chat_message_history


    def run(self, query: str):
        emotion = self.emotion_chain(query)
        print("当前用户的情绪:", self.CURRENT_EMOTION)
        print("当前设定:", self.MOODS[self.CURRENT_EMOTION]["roleSet"])
        print("prompt1:", self.prompt.messages)
        self.prompt = ChatPromptTemplate.from_messages(
            [
                (
                    "system",
                    # self.SYSTEMPL,
                    self.SYSTEMPL.format(who_you_are=self.MOODS[self.CURRENT_EMOTION]["roleSet"]),
                ),
                # MessagesPlaceholder(variable_name=self.MEMORY_KEY),
                (
                    "user",
                    "{query}",
                ),
                MessagesPlaceholder(variable_name="agent_scratchpad"),
            ]
        )
        print("prompt2:", self.prompt.messages)
        result = self.agent_executor.invoke({"query": query, "chat_history": self.memory.messages})
        # result = self.agent_executor.invoke({"query": query})
        # return self.agent_executor.get_response()
        return result

    def emotion_chain(self, query: str):
        prompt_emotion = """根据用户的输入判断用户的情绪,回应规则对照下面：
        1. 内容为负面情绪，只返回"depressed"，不要有其他内容，例如压抑、抑郁的语句.
        2. 内容为正面情绪，只返回"friendly"，不要有其他内容，例如友好的、礼貌的语句.
        3. 内容为中性情绪，只返回"default"，不要有其他内容.
        4. 内容为愤怒生气情绪的内容，只返回"angry"，不要有其他内容，例如愤怒、辱骂、笨蛋、仇恨的语句.
        5. 内容包含情绪十分开心，只返回"cheerful"，不要有其他内容，例如高兴的、狂喜的、兴奋、称赞的语句.
        用户输入内容:{query}"""
        chat = ChatPromptTemplate.from_template(prompt_emotion) | self.chatmodel | StrOutputParser()
        result = chat.invoke({"query": query})
        self.CURRENT_EMOTION = result
        return result


@app.get("/")
def read_root():
    return {"message": "Hello, World!"}


@app.post("/chat")
def chat(query: str):
    master = Master()
    return master.run(query)


@app.post("/add_urls")
def add_urls(url: str):
    loader = WebBaseLoader(url)
    docs = loader.load()
    documents = RecursiveCharacterTextSplitter(
        chunk_size=800,
        chunk_overlap=50
    ).split_documents(docs)
    # 引入向量数据库
    qdrant = Qdrant.from_documents(
        documents,
        DashScopeEmbeddings(model="text-embedding-v1"),
        path="",
        collection_name="yunshi_2024"
    )
    print("向量数据库创建完成")
    print(qdrant.co)
    return {"ok": "添加成功"}


@app.post("/add_pdfs")
def add_pdfs():
    return {"message": "PDFs added successfully"}


@app.post("/add_texts")
def add_text():
    return {"message": "Text added successfully"}


@app.websocket("/ws")
async def websocket_endpoint(webSocket: WebSocket):
    await webSocket.accept()
    try:
        while True:
            data = await webSocket.receive_text()
            await webSocket.send_text(f"Message text was: {data}")
    except WebSocketDisconnect:
        print("WebSocket connection closed")
        await webSocket.close()


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8001)